import dolphindb as ddb
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
import seaborn as sns
import matplotlib.dates as mdates
from scipy.stats import zscore, gaussian_kde
import pandas as pd

#连接dolphin DB
s = ddb.session()# 创建会话
s.connect(host='101.95.132.98',port=33147, userid='CW', password='123456') # 建立连接

t1 =  'select * from atm_ts  '#从DB 中获取原始数据
data_new = s.run(t1)

data_new['hour'] = pd.to_datetime(data_new['hour'], format='%Y.%m.%dT%H')

# Calculate Z-score
data_new['z_score'] = data_new.groupby('code')['hour'].transform(lambda x: zscore(x, nan_policy='omit'))

# Filter the data based on Z-score (-6, 6)
data_for_densityplot_new_filtered = data_new[np.abs(data_new['z_score']) < 6]

Identify the top 10 codes with most data points
data_count_by_code = data_new.groupby('code').size().reset_index(name='count')
top_4_codes = data_count_by_code.nlargest(4, 'count')['code'].tolist()

Filter the data for plotting the last 5 days
data_for_lineplot_new = data_new[data_new['hour'] > data_new['hour'].max() - pd.Timedelta(days=5)]



def plot_all_iv_diff_and_density_for_codes_adjusted_xaxis_fixed(codes, data_for_lineplot, data_for_densityplot):
    n = len(codes)
    gs = gridspec.GridSpec(n, 2, width_ratios=[3, 1], hspace=0.5)
    fig = plt.figure(figsize=(15, 10*n))
    
    for i, code in enumerate(codes):
        ax1 = fig.add_subplot(gs[i, 0])
        
        ax2 = fig.add_subplot(gs[i, 1])

        data_for_lineplot_code = data_for_lineplot[data_for_lineplot['code'] == code].copy()
        data_for_densityplot_code = data_for_densityplot[data_for_densityplot['code'] == code].copy()
        data_for_densityplot_code['z_score'] = zscore(data_for_densityplot_code['diff'], nan_policy='omit')
        data_for_densityplot_code = data_for_densityplot_code[np.abs(data_for_densityplot_code['z_score']) < 6]

        # Lineplot
        lineplot = sns.lineplot(data=data_for_lineplot_code, x='hour', y='diff', ax=ax1)
        
        # Plotting the mean of skewness on the lineplot
        mean_diff = data_for_lineplot_code['diff'].mean()
        ax1.axhline(mean_diff, color='red', linestyle='--')
        
        ax1.set_title('')
        ax1.set_xlabel('') #去x轴标题
        ax1.set_ylabel(f' {code}')
        
        # Adjust the x-axis date format
        if i == len(codes) - 1:
            ax1.xaxis.set_major_locator(mdates.DayLocator())
            ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
            ax1.tick_params(axis='x', rotation=45)    
        else:
            ax1.set_xticks([])

        # Density plot
        if len(data_for_densityplot_code) > 1:  # Only plot density if there are more than 1 data point
            density = gaussian_kde(data_for_densityplot_code['diff'])
            x = np.linspace(min(data_for_densityplot_code['diff']), max(data_for_densityplot_code['diff']), 500)
            y = density(x)
            line_color = lineplot.get_lines()[0].get_color()
            ax2.fill_betweenx(x, y, color=line_color, alpha=0.5)
        
        # Plotting the mean of skewness on the density plot
        ax2.axhline(mean_diff, color='red', linestyle='--')
        
        if i == len(codes) - 1:
            ax2.set_xlabel('Density')
        else:
            ax2.set_title('')
            ax2.set_ylabel('')
            ax2.set_yticks([])
                            
    plt.tight_layout()
    plt.show()

# Filter out codes with less than 2 data points for density plotting
codes_with_enough_data = data_for_densityplot_new_filtered['code'].value_counts()
codes_with_enough_data = codes_with_enough_data[codes_with_enough_data > 1].index.tolist()

# Filter the data based on the codes with enough data
filtered_data_for_lineplot = data_for_lineplot_new[data_for_lineplot_new['code'].isin(codes_with_enough_data)]
filtered_data_for_densityplot = data_for_densityplot_new_filtered[data_for_densityplot_new_filtered['code'].isin(codes_with_enough_data)]

# Plot using the adjusted function
plot_all_iv_diff_and_density_for_codes_adjusted_xaxis_fixed(codes_with_enough_data, filtered_data_for_lineplot, filtered_data_for_densityplot)

# Call the function to plot with adjusted x-axis and fixed density plots
# plot_all_iv_diff_and_density_for_codes_adjusted_xaxis_fixed(top_4_codes, data_for_lineplot_new, data_for_densityplot_new_filtered)

























































































